import re

import pandas as pd

from iflytech_assistant.assistant.dataclasses import RagData

user_prompt = """
我是一个帮你优化日常对话的小助手。我的主要任务是帮你把话说得更自然、更有吸引力，但还是保持你想表达的意思。

我在跟我的{target}聊天，我想给对方表达下面这段话，但是不知道怎么说才好。
```
{user_input}
```
麻烦帮我按照{preferenced_style}的风格改写一下呗。
直接输出给我结果就行了，不需要解释。
"""


df = pd.read_excel(".vscode/上屏发送数据.xlsx")
# remove text in "用户偏好" column except English words and Chinese characters
df["用户偏好"] = df["用户偏好"].apply(
    lambda x: re.sub(r"[^\u4e00-\u9fa5a-zA-Z]", "", x)
)
df["i_clickword"] = df["i_clickword"].apply(lambda x: re.sub(r"A、.*$", "", x))

df = df[["d_submodeid", "content", "userinput", "i_clickword", "send_word", "用户偏好"]]
df["d_submodeid"] = df["d_submodeid"].apply(lambda x: x.replace("发", ""))


# group by "d_submodeid" and "用户偏好" and userinput

df = (
    df.groupby(["d_submodeid", "用户偏好", "userinput"])
    .agg(lambda x: "\n".join(x))
    .reset_index()
)

final_df = []
for i, row in df.iterrows():
    data: RagData = RagData(
        input=row["userinput"],
        target=row["d_submodeid"],
        tag=row["用户偏好"],
        mode="polish",
        examples=row["i_clickword"].split("\n"),
    )

    prompt = user_prompt.format(
        user_input=data.input,
        target=data.target,
        preferenced_style=data.tag,
    )

    final_df.append(
        {
            "input": prompt,
            "target": data.examples[0],
        }
    )

final_df = pd.DataFrame(final_df)
# shuffle
final_df = final_df.sample(frac=1).reset_index(drop=True)
# final_df.to_json(
#     ".vscode/finetune.jsonl", orient="records", lines=True, force_ascii=False
# ) final_df.to_excel(".vscode/finetune.xlsx", index=False)

